package com.ruijie.autotest.log.common.crypto;

import com.minelx.landisk.filesystem.core.File;
import com.minelx.util.IOUtil;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import java.io.ByteArrayInputStream;
import java.io.OutputStream;
import java.nio.file.StandardOpenOption;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

import static com.minelx.landisk.filesystem.core.PathFactory.file;
import static com.minelx.util.IOUtil.bytesOf;
import static javax.crypto.Cipher.DECRYPT_MODE;
import static javax.crypto.Cipher.ENCRYPT_MODE;

public class AES {
	public static final String PKCS_5_PADDING = "PKCS5Padding";
	public static final String AES_CBC_PADDING = "AES/CBC/"+ PKCS_5_PADDING;

	public static byte[] encrypt(byte[] rawBytes, String password, String ivPassword) throws Exception {
		Cipher result = create(password, ivPassword, ENCRYPT_MODE);

		return Base64.getEncoder().encode(result.doFinal(rawBytes));
	}

	public static byte[] decrypt(byte[] rawBytes, String password, String ivPassword)
			throws Exception {
		Cipher result = create(password, ivPassword, DECRYPT_MODE);
		return result.doFinal(Base64.getDecoder().decode(new String(rawBytes)));
	}

	public static Cipher create(String password, String ivPassword, int decryptMode) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, InvalidAlgorithmParameterException {
		Cipher result = Cipher.getInstance(AES_CBC_PADDING);
		result.init(decryptMode,
				secretKeyOf(password),
				new IvParameterSpec(ivPassword.getBytes())
		);
		return result;
	}

	private static SecretKey secretKeyOf(String password) throws NoSuchAlgorithmException {
		byte[] passwordAsBytes = password.getBytes();
		KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
		SecureRandom secureRandom = SecureRandom.getInstance("SHA1PRNG");
		secureRandom.setSeed(passwordAsBytes);
		keyGenerator.init(256, secureRandom);
		return keyGenerator.generateKey();
	}

	public static List<File> decryptedFiles(List<File> decryptedFiles, Crypto decryption) throws Exception {
		List<File> result = new ArrayList<>();
		for (File decryptedFile : decryptedFiles) {
			System.out.println("decryptedFile = " + decryptedFile.getName());

			byte[] decryptedAsBytes = decryption.doCrypto(bytesOf(decryptedFile.inputStream()));
			File dest = decryptedFile.sibling("decrypted_" + decryptedFile.getName()).as(file());
			try (OutputStream outputStream = dest.outputStream(StandardOpenOption.CREATE)) {
				IOUtil.exchange(new ByteArrayInputStream(decryptedAsBytes), outputStream);
			}
			result.add(dest);
		}
		return result;
	}
}
